import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn

max_epoch = 100
batch_size = 1
learning_rate = 0.1

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

xs = np.array([1, 2, 3])
ys = np.array([2, 4, 6])

# prepare data
class LinearRegressionDataset(Dataset):
    def __init__(self, xs, ys):
        super(LinearRegressionDataset, self).__init__()
        self.xs = xs
        self.ys = ys
        # self.xs = torch.tensor(xs, dtype=torch.float32)
        # self.ys = torch.tensor(ys, dtype=torch.float32)

    def __getitem__(self, index):
        x = self.xs[index]
        y = self.ys[index]
        x = torch.tensor(x, dtype=torch.float32)
        y = torch.tensor(y, dtype=torch.float32)
        return x, y

    def __len__(self):
        return len(self.xs)

def test_dataset():
    dataset = LinearRegressionDataset(xs, ys)
    dataloader = DataLoader(
                    dataset,
                    batch_size=batch_size,
                    shuffle=True
                )
    for step, batch in enumerate(dataloader):
        x, y = batch
        print(x, y)

# model
class LinearRegressionModel(nn.Module):
    def __init__(self):
        super(LinearRegressionModel, self).__init__()
        self.fc = nn.Linear(1, 1) # y = w * x + b

    def forward(self, x):
        x = self.fc(x)
        return x


# train
def train(dataloader, model):
    # loss function and optimizer
    criterion = nn.MSELoss()  # loss function
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    # w = w - alpha * dL/dw
    # b = b - alpha * dL/db

    for epoch in range(max_epoch):
        for step, batch in enumerate(dataloader):
            batch_x, batch_y = batch
            batch_x = batch_x.to(device)
            batch_y = batch_y.to(device)

            pred_y = model(batch_x)
            loss = criterion(pred_y, batch_y)

            # update parameter
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if step % 10 == 0:
                print("Epoch: {:d}, Step: {:d}, Loss: {:.4f}".format(epoch, step, loss.item()))

if __name__ == '__main__':
    # test_dataset()
    dataset = LinearRegressionDataset(xs, ys)
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
    )

    model = LinearRegressionModel()
    model.to(device)

    train(dataloader, model)

    test_x = torch.tensor([5], dtype=torch.float32)
    test_y = model(test_x)
    print(test_y)

    print(model.fc.weight[0][0])
    print(model.fc.bias)











